import requests
from requests.auth import HTTPBasicAuth
import time
import random
import functools
from datetime import datetime
import argparse
import os
import pandas as pd
import winsound
#  版本说明：提交前进行检查，确认要有结果，才可以提交，会出现长时间没有状态返回
# 创建命令行参数解析器
parser = argparse.ArgumentParser(description='WorldQuant Alpha提交器')
parser.add_argument('--credentials_file', type=str, default="credentials.txt", help='手动创建账号文件，文件首行账号，第二行密码,不需要其他符号')
parser.add_argument('--start_date', type=str, default="05-01", help='开始日期 (MM-DD格式)')
parser.add_argument('--end_date', type=str, default="07-12", help='结束日期 (MM-DD格式)')
parser.add_argument('--alpha_num', type=int, default=5000, help='要检查的Alpha数量')
parser.add_argument('--sharpe_th', type=float, default=1.25, help='Sharpe阈值')
parser.add_argument('--fitness_th', type=float, default=1.0, help='Fitness阈值')
parser.add_argument('--turnover_th', type=float, default=0.3, help='Turnover阈值')
parser.add_argument('--submit_delay', type=int, default=70, help='提交之间的延迟时间(秒)')
parser.add_argument('--max_submitted_change', type=int, default=2, help='最大允许的已提交Alpha变化数量')
parser.add_argument('--region', type=str, default="USA", help='地区')
parser.add_argument('--blacklist_file', type=str, default="blacklist.txt", help='黑名单文件路径')

args = parser.parse_args()
condition = True  # 声音开关


# 从文件读取凭据
def read_credentials(file_path):
    username = ""
    password = ""
    try:
        if os.path.exists(file_path):
            with open(file_path, 'r') as file:
                lines = file.readlines()
                if len(lines) >= 1:
                    username = lines[0].strip()
                if len(lines) >= 2:
                    password = lines[1].strip()
            return username, password
        else:
            print(f"凭据文件 {file_path} 不存在")
            return "", ""
    except Exception as e:
        print(f"读取凭据文件时出错: {e}")
        return "", ""


# 读取黑名单（文件不存在时创建）
def read_blacklist(file_path):
    blacklist = set()
    try:
        if not os.path.exists(file_path):
            with open(file_path, 'w') as file:
                pass
            print(f"黑名单文件 {file_path} 不存在，已创建新文件")
        else:
            with open(file_path, 'r') as file:
                for line in file:
                    blacklist.add(line.strip())
            print(f"已从黑名单文件中读取 {len(blacklist)} 个Alpha ID")
    except Exception as e:
        print(f"读取或创建黑名单文件时出错: {e}")
    return blacklist


# 更新黑名单（实时写入）
def update_blacklist(file_path, alpha_id):
    try:
        with open(file_path, 'a') as file:
            file.write(f"{alpha_id}\n")
        print(f"已实时将失败的Alpha ID {alpha_id} 添加到黑名单")
        return True
    except Exception as e:
        print(f"实时更新黑名单文件时出错: {e}")
        return False


# 检查Alpha提交状态（带超时）
def get_check_submission(s, alpha_id):
    try:
        while True:
            result = s.get(f"https://api.worldquantbrain.com/alphas/{alpha_id}/check", timeout=30)
            if "retry-after" in result.headers:
                time.sleep(float(result.headers["Retry-After"]))
            else:
                break
        if result.json().get("is", 0) == 0:
            print(f"Alpha {alpha_id}: logged out，返回 'sleep'")
            return "sleep"
        checks_df = pd.DataFrame(result.json()["is"]["checks"])
        # 检查 SELF_CORRELATION 是否为 "nan"
        self_correlation_value = checks_df[checks_df["name"] == "SELF_CORRELATION"]["value"].values[0]
        pc = self_correlation_value
        if any(checks_df["result"] == "ERROR"):
            print(f"Alpha {alpha_id}: \033[31m ERROR \033[0m，检查失败")
            return "ERROR"
        if any(checks_df["result"] == "FAIL"):
            print(f"Alpha {alpha_id}: \033[31m FAIL \033[0m，检查失败")
            return "FAIL"
        if pd.isna(self_correlation_value) or str(self_correlation_value).lower() == "nan":
            print(f"Alpha {alpha_id}: SELF_CORRELATION 为 \033[31m nan \033[0m，检查失败")
            return "nan"
        else:
            print(f"\033[34m  Alpha {alpha_id}: 检查通过  \033[0m ")
            return pc
    except requests.exceptions.Timeout:
        print(f"Alpha {alpha_id}: 检查超时（30秒未返回结果）")
        return "timeout"
    except Exception as e:
        print(f"catch: {alpha_id} - {str(e)}")
        return "error"


# 提交Alpha
def submit_alpha(s, alpha_id):
    max_retries = 3
    retry_delay = 20
    status_code = None
    for retry in range(max_retries):
        if retry > 0:
            print(f"连接问题，等待 {retry_delay} 秒后第 {retry + 1} 次尝试...")
            time.sleep(retry_delay)
        try:
            response = s.post(
                f"https://api.worldquantbrain.com/alphas/{alpha_id}/submit",
                json={},
                timeout=(10, 30)
            )
            status_code = response.status_code
            print(f"提交状态码: {status_code}")
            if status_code < 300:
                return True, status_code  # 成功提交
            elif status_code == 400:
                print(f"Alpha {alpha_id}: 状态码 400 (Bad Request)，提交失败")
                return False, status_code  # 返回失败，不触发黑名单
            elif status_code == 403:
                print(f"Alpha {alpha_id}: 状态码 403 (Forbidden)，提交失败")
                return False, status_code  # 返回失败，触发黑名单
            elif status_code == 429:
                print(f"触发速率限制，等待更长时间...")
                time.sleep(retry_delay * 2)
                continue
        except requests.exceptions.Timeout:
            print("请求超时，将重试...")
            continue
        except requests.exceptions.ConnectionError as e:
            print(f"连接错误: {str(e)}")
            continue
        except Exception as e:
            print(f"提交出错: {str(e)}")
    return False, status_code  # 重试失败后返回


# 读取凭据
username, password = read_credentials(args.credentials_file)
if not username or not password:
    print("未能获取有效的用户名或密码，请检查凭据文件")
    exit()

# 读取黑名单
blacklist = read_blacklist(args.blacklist_file)

# 设置其他参数
sharpe_th = args.sharpe_th
fitness_th = args.fitness_th
turnover_th = args.turnover_th
SUBMIT_DELAY = args.submit_delay
MAX_SUBMITTED_CHANGE = args.max_submitted_change
start_date = args.start_date
end_date = args.end_date
alpha_num = args.alpha_num
region = args.region

# 创建会话
session = requests.Session()
session.auth = HTTPBasicAuth(username, password)
session.headers.update({
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36',
    'Accept': 'application/json',
    'Content-Type': 'application/json'
})
session.request = functools.partial(session.request, timeout=(10, 30))


# 认证并获取用户ID
def login():
    response = session.post('https://api.worldquantbrain.com/authentication')
    if response.status_code < 300:
        auth_data = response.json()
        user_id = auth_data['user']['id']
        if 'token' in auth_data:
            session.headers.update({'Authorization': f'Bearer {auth_data["token"]}'})
        return session
    else:
        print(f"认证失败: {response.status_code}")
        return None


# 获取特定状态的Alpha数量
def get_alpha_count(status):
    s = login()
    if not s:
        return None
    try:
        url = f"https://api.worldquantbrain.com/users/self/alphas?limit=1&status={status}"
        response = s.get(url)
        if response.status_code < 300:
            count = response.json().get('count', 0)
            return count
        else:
            print(f"获取状态为 '{status}' 的Alpha数量失败: {response.status_code}")
            return None
    except Exception as e:
        print(f"获取状态为 '{status}' 的Alpha数量时出错: {e}")
        return None


# 获取有效Alpha
def get_alphas(start_date, end_date, sharpe_th, fitness_th, turnover_th, region, alpha_num, usage):
    s = login()
    output = []
    count = 0
    current_year = datetime.now().strftime('%Y')
    for i in range(0, alpha_num, 100):
        print(i)
        # 修改URL，添加fitness上限条件 &is.fitness%3C2.5
        url_e = f"https://api.worldquantbrain.com/users/self/alphas?limit=100&offset={i}" \
                f"&status=UNSUBMITTED%1FIS_FAIL&dateCreated%3E={current_year}-{start_date}" \
                f"T00:00:00-04:00&dateCreated%3C{current_year}-{end_date}" \
                f"T00:00:00-04:00&is.fitness%3E{fitness_th}&is.fitness%3C2.5&is.sharpe%3E{sharpe_th}" \
                f"&settings.region={region}&order=is.sharpe&hidden=false&type!=SUPER" \
                f"&is.turnover%3C{turnover_th}"
        # 对于负值，使用 &is.fitness%3E-2.5 作为下限
        url_c = f"https://api.worldquantbrain.com/users/self/alphas?limit=100&offset={i}" \
                f"&status=UNSUBMITTED%1FIS_FAIL&dateCreated%3E={current_year}-{start_date}" \
                f"T00:00:00-04:00&dateCreated%3C{current_year}-{end_date}" \
                f"T00:00:00-04:00&is.fitness%3C-{fitness_th}&is.fitness%3E-2.5&is.sharpe%3C-{sharpe_th}" \
                f"&settings.region={region}&order=is.sharpe&hidden=false&type!=SUPER" \
                f"&is.turnover%3C{turnover_th}"
        urls = [url_e]
        if usage != "submit":
            urls.append(url_c)
        for url in urls:
            response = s.get(url)
            try:
                alpha_list = response.json()["results"]
                for j in range(len(alpha_list)):
                    alpha_id = alpha_list[j]["id"]
                    if alpha_id in blacklist:
                        print(f"跳过ID为 {alpha_id} 的Alpha，因为它在黑名单中")
                        continue
                    name = alpha_list[j]["name"]
                    dateCreated = alpha_list[j]["dateCreated"]
                    sharpe = alpha_list[j]["is"]["sharpe"]
                    fitness = alpha_list[j]["is"]["fitness"]

                    # # 添加额外的检查，确保fitness绝对值不超过2.0
                    # if abs(fitness) > 10.0:
                    #     print(f"跳过ID为 {alpha_id} 的Alpha，因为fitness值为 {fitness}，超过2.0")
                    #     continue

                    turnover = alpha_list[j]["is"]["turnover"]
                    margin = alpha_list[j]["is"]["margin"]
                    longCount = alpha_list[j]["is"]["longCount"]
                    shortCount = alpha_list[j]["is"]["shortCount"]
                    decay = alpha_list[j]["settings"]["decay"]
                    exp = alpha_list[j]['regular']['code']
                    count += 1
                    checks = alpha_list[j].get("is", {}).get("checks", [])
                    has_failed_checks = any(check.get('result') == 'FAIL' for check in checks if check)
                    if has_failed_checks:
                        print(f"跳过ID为 {alpha_id} 的Alpha，因为它有失败的检查项")
                        continue
                    if (longCount + shortCount) > 100 and turnover < turnover_th:
                        if sharpe < -sharpe_th:
                            exp = "-%s" % exp
                        rec = [alpha_id, exp, sharpe, turnover, fitness, margin, dateCreated, decay]
                        print(rec)
                        if turnover > 0.25:
                            rec.append(decay + 2)
                        elif turnover > 0.2:
                            rec.append(decay)
                        elif turnover > 0.15:
                            rec.append(decay)
                        else:
                            rec.append(decay)
                        output.append(rec)
            except Exception as e:
                print(f"{i} finished re-login: {e}")
                s = login()
    print("count: %d" % count)
    return output


# 主程序
def main():
    print("=== WorldQuant Alpha提交器 ===")
    print(f"开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"凭据文件: {args.credentials_file}")
    print(f"黑名单文件: {args.blacklist_file}")
    print(f"日期范围: {start_date} 至 {end_date}")
    print(f"检查的Alpha数量: {alpha_num}")
    print(f"地区: {region}")
    print(f"Sharpe阈值: {sharpe_th}")
    print(f"Fitness阈值: {fitness_th}")
    print(f"Turnover阈值: {turnover_th}")
    print(f"提交延迟: {SUBMIT_DELAY}秒")
    print(f"最大允许已提交Alpha变化数量: {MAX_SUBMITTED_CHANGE}")

    s = login()
    if not s:
        print("登录失败，程序退出")
        return

    initial_submitted_count = get_alpha_count("ACTIVE")
    if initial_submitted_count is not None:
        print(f"平台上已提交的Alpha数量: {initial_submitted_count}")
    else:
        print("无法计算ACTIVE，请检查登录凭据或网络连接")

    print("\n正在获取Alpha列表...")
    print(f"\n搜索符合条件的有效alpha (Sharpe >= {sharpe_th}, Fitness >= {fitness_th}, Turnover < {turnover_th})...")
    valid_alphas_data = get_alphas(start_date, end_date, sharpe_th, fitness_th, turnover_th, region, alpha_num,
                                   "submit")
    valid_alphas = [alpha[0] for alpha in valid_alphas_data]
    alpha_metrics = {alpha[0]: {"sharpe": alpha[2], "fitness": alpha[4]} for alpha in valid_alphas_data}
    print(f"找到 {len(valid_alphas)} 个有效Alpha（不包含失败检查项和黑名单中的Alpha）")

    if not valid_alphas:
        print("没有发现符合条件的有效Alpha，无需提交。")
        return

    print(f"\n准备自动提交 {len(valid_alphas)} 个有效Alpha")
    submitted = 0
    failed = 0

    # 第一轮提交
    for i, alpha_id in enumerate(valid_alphas):
        print(f"检查 {i + 1}/{len(valid_alphas)}: {alpha_id}  [Sharpe: {alpha_metrics[alpha_id]['sharpe']}, Fitness: {alpha_metrics[alpha_id]['fitness']}]" )

        # 先检查Alpha状态，处理 "sleep" 和超时逻辑
        retry_count = 0
        while True:
            check_result = get_check_submission(s, alpha_id)
            if check_result == "sleep":
                print(f"检查结果: sleep，重新登录并等待 40 秒")
                s = login()
                if not s:
                    print("重新登录失败，跳过此 Alpha")
                    failed += 1
                    if update_blacklist(args.blacklist_file, alpha_id):
                        blacklist.add(alpha_id)
                    break
                time.sleep(40)
                continue
            elif check_result == "timeout" and retry_count < 1:
                print(f"检查超时，第 {retry_count + 1} 次重试")
                retry_count += 1
                continue
            else:
                break

        #
        if check_result == "timeout":
            print(f"检查结果: 第二次超时，直接提交 Alpha {alpha_id}")
        # 根据检查结果处理
        elif check_result == "nan":
            print(f"检查结果: 失败 (\033[31m nan \033[0m)，列入黑名单")
            failed += 1
            if update_blacklist(args.blacklist_file, alpha_id):
                blacklist.add(alpha_id)
            continue
        elif check_result == "ERROR":
            print(f"检查结果: 失败 (ERROR)，列入黑名单")
            failed += 1
            if update_blacklist(args.blacklist_file, alpha_id):
                blacklist.add(alpha_id)
            continue
        elif check_result == "FAIL":
            print(f"检查结果: 错误 (FAIL)，列入黑名单")
            failed += 1
            if update_blacklist(args.blacklist_file, alpha_id):
                blacklist.add(alpha_id)
            continue
        else:
            print(f"检查结果: 通过 (SELF_CORRELATION: {check_result})，开始提交")


        success, status_code = submit_alpha(s, alpha_id)
        if success:
            print(f"提交结果: 已提交! 状态码: {status_code}")
            if status_code == 201:
                if condition:
                    # 播放系统默认声音
                    winsound.MessageBeep()
                    # # winsound.Beep(frequency, duration)  # 频率(Hz)和持续时间(ms)
                    winsound.Beep(1000, 500)  # 1000Hz, 500毫秒
            submitted += 1
            delay = SUBMIT_DELAY + random.uniform(5, 15)
            print(f"等待 {delay:.2f} 秒...")
            time.sleep(delay)
            current_submitted_count = get_alpha_count("ACTIVE")
            if current_submitted_count is None:
                print("无法获取当前已提交Alpha数量，继续执行...")
            else:
                change = abs(current_submitted_count - initial_submitted_count)
                print(f"总共成功提交数量： {change} !")
                if change >= MAX_SUBMITTED_CHANGE:
                    print(f"警告：已提交的Alpha数量变化 ({change}) 超过阈值 ({MAX_SUBMITTED_CHANGE})!")
                    print(f"预期已提交数量: {submitted}, 实际已提交数量: {change}")
                    print("程序停止执行。")
                    return
        else:
            print(f"提交结果: 失败! 状态码: {status_code}")
            failed += 1
            if status_code not in (400, 429):
              if update_blacklist(args.blacklist_file, alpha_id):
                blacklist.add(alpha_id)

    print(f"\n第一轮提交:")
    print(f"总共: {len(valid_alphas)} 个Alpha")
    print(f"已提交: {submitted} 个")
    print(f"失败: {failed} 个")

    # 第二轮重试
    if failed > 0:
        print(f"\n开始重新检查和提交失败的Alpha")
        retry_submitted = 0
        retry_failed = 0
        valid_alphas_data = get_alphas(start_date, end_date, sharpe_th, fitness_th, turnover_th, region, alpha_num,"submit")
        valid_alphas = [alpha[0] for alpha in valid_alphas_data if alpha[0] not in blacklist]

        for i, alpha_id in enumerate(valid_alphas):
            print(f"重新检查 {i + 1}/{len(valid_alphas)}: {alpha_id}")

            # 先检查Alpha状态，处理 "sleep" 和超时逻辑
            retry_count = 0
            while True:
                check_result = get_check_submission(s, alpha_id)
                if check_result == "sleep":
                    print(f"检查结果: sleep，重新登录并等待 40 秒")
                    s = login()
                    if not s:
                        print("重新登录失败，跳过此 Alpha")
                        retry_failed += 1
                        if update_blacklist(args.blacklist_file, alpha_id):
                            blacklist.add(alpha_id)
                        break
                    time.sleep(40)
                    continue
                elif check_result == "timeout" and retry_count < 1:
                    print(f"检查超时，第 {retry_count + 1} 次重试")
                    retry_count += 1
                    continue
                else:
                    break

            # 如果第二次检查仍超时，直接提交
            if check_result == "timeout":
                print(f"检查结果: 第二次超时，直接提交 Alpha {alpha_id}")
            # 根据检查结果处理
            elif check_result == "nan":
                print(f"检查结果: 失败 (nan)，列入黑名单")
                retry_failed += 1
                if update_blacklist(args.blacklist_file, alpha_id):
                    blacklist.add(alpha_id)
                continue
            elif check_result == "ERROR":
                print(f"检查结果: 错误 (ERROR)，列入黑名单")
                retry_failed += 1
                if update_blacklist(args.blacklist_file, alpha_id):
                    blacklist.add(alpha_id)
                continue
            elif check_result == "FAIL":
                print(f"检查结果: 错误 (FAIL)，列入黑名单")
                retry_failed += 1
                if update_blacklist(args.blacklist_file, alpha_id):
                    blacklist.add(alpha_id)
                continue
            else:
                print(f"检查结果: 通过 (SELF_CORRELATION: {check_result})，开始提交")

            # 检查通过或超时后提交
            current_submitted_count = get_alpha_count("ACTIVE")
            if current_submitted_count is None:
                print("无法获取当前已提交Alpha数量，继续执行...")
            else:
                change = abs(current_submitted_count - initial_submitted_count)
                print(f"总共成功提交数量： {change} !")
                if change >= MAX_SUBMITTED_CHANGE:
                    print(f"警告：已提交的Alpha数量变化 ({change}) 超过阈值 ({MAX_SUBMITTED_CHANGE})!")
                    print(f"预期已提交数量: {submitted}, 实际已提交数量: {change}")
                    print("程序停止执行。")
                    return

            success, status_code = submit_alpha(s, alpha_id)
            if success:
                print(f"提交结果: 已提交! 状态码: {status_code}")
                retry_submitted += 1
                submitted += 1
                failed -= 1
                delay = SUBMIT_DELAY + random.uniform(5, 15)
                print(f"等待 {delay:.2f} 秒...")
                time.sleep(delay)
            else:
                print(f"提交结果: 失败! 状态码: {status_code}")
                retry_failed += 1
                if status_code not in (400, 429):
                    if update_blacklist(args.blacklist_file, alpha_id):
                        blacklist.add(alpha_id)

        print(f"\n第二轮提交:")
        print(f"尝试重提交: {len(valid_alphas)} 个")
        print(f"重提交成功: {retry_submitted} 个")
        print(f"重提交失败: {retry_failed} 个")

    # 最终总结
    print(f"\n最终结果:")
    print(f"总共: {len(valid_alphas)} 个Alpha")
    print(f"已提交: {submitted} 个")
    print(f"失败: {failed} 个")
    print(f"已提交率: {(submitted / len(valid_alphas) * 100):.2f}%")
    print(f"完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

    final_submitted_count = get_alpha_count("ACTIVE")
    if final_submitted_count is not None:
        print(f"本次运行成功新增提交: {final_submitted_count - initial_submitted_count}")
        print(f"本次运行已提交: {submitted}")


if __name__ == "__main__":
    main()